package com.zmh.wt.healthmonitor.utils;

import com.zmh.wt.healthmonitor.entity.MyComplex;
import com.zmh.wt.healthmonitor.entity.SpectrumDo;
import com.zmh.wt.healthmonitor.entity.TxtFileDo;
import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;
import org.junit.jupiter.api.Test;

import java.util.stream.Stream;

/**
 * @Author MH.Zhang
 * @Classname FFTUtil
 * @Date 2023/3/28 16:07
 */
public class FFTUtil {
    public static void main(String[] args) {
        TxtFileDo txtFileDo = TxtUtil.readTxtByPath("E:\\txt测试数据\\A01-20190722131913.txt");
        double interval = 1;
        Double[] orgData = txtFileDo.getData();
        Double[] data = TxtUtil.highpass(orgData).getData();

        int N = 2;
        int length = data.length;
        while (N < data.length) {
            N *= 2;
        }

        MyComplex[] input = new MyComplex[N];
        for (int i = 0; i <= N - 1; i++) {
            if (i < length) {
                input[i] = new MyComplex(data[i], 0);
            } else {
                input[i] = new MyComplex(0, 0);
            }
        }

        input = getFFT(input, N);//傅里叶变换
        Double[] modArray = MyComplex.toModArray(input);//计算傅里叶变换得到的复数数组的模值
        Double[] norArray = new Double[modArray.length];
        for (int i = 0; i <= N - 1; i++) {
            //的模值数组除以N再乘以2
            norArray[i] = modArray[i] / N * 2;
        }


//        double[] doubles = Stream.of(txtFileDo.getData()).mapToDouble(Double::doubleValue).toArray();
//        SpectrumDo spectrumDo = caculateSpectrum(doubles, interval, true, flyLen);
//        double[] amplitude = spectrumDo.getAmplitude();
//        TxtUtil.writeTxtFile("E:\\txt测试数据\\spectrum.txt", norArray);

    }

    @Test
    public void testFFT() {
        TxtFileDo txtFileDo = TxtUtil.readTxtByPath("E:\\txt测试数据\\A01-20190722131913.txt");
        double interval = 44100.0;
        Double[] orgData = txtFileDo.getData();
        Double[] data = TxtUtil.highpass(orgData).getData();

        int N = 2;
        while (N < data.length) {
            N *= 2;
        }

        double[] doubles = Stream.of(data).mapToDouble(Double::doubleValue).toArray();
        SpectrumDo spectrumDo = caculateSpectrum(doubles, interval, true, N);
        double[] amplitude = spectrumDo.getAmplitude();
        TxtUtil.writeTxtFile("E:\\txt测试数据\\spectrum.txt", amplitude);
    }


    public static MyComplex[] getFFT(MyComplex[] input, int N) {
        if ((N / 2) % 2 == 0) {
            MyComplex[] even = new MyComplex[N / 2];// 偶数
            MyComplex[] odd = new MyComplex[N / 2];// 奇数
            for (int i = 0; i < N / 2; i++) {
                even[i] = input[2 * i];
                odd[i] = input[2 * i + 1];
            }
            even = getFFT(even, N / 2);
            odd = getFFT(odd, N / 2);
            for (int i = 0; i < N / 2; i++) {
                MyComplex[] res = MyComplex.butterfly(even[i], odd[i], MyComplex.GetW(i, N));
                input[i] = res[0];
                input[i + N / 2] = res[1];
            }
            return input;
        } else {// 两点DFT,直接进行碟形运算
            MyComplex[] res = MyComplex.butterfly(input[0], input[1], MyComplex.GetW(0, N));
            return res;
        }
    }


    public static Complex[] fft(Complex[] x) {
        int N = x.length;

        // base case
        if (N == 1) {
            return new Complex[]{x[0]};
        }

        // radix 2 Cooley-Tukey FFT
        if (N % 2 != 0) {
            throw new RuntimeException("N is not a power of 2");
        }

        // fft of even terms
        Complex[] even = new Complex[N / 2];
        for (int k = 0; k < N / 2; k++) {
            even[k] = x[2 * k];
        }
        Complex[] q = fft(even);

        // fft of odd terms
        Complex[] odd = even;  // reuse the array
        for (int k = 0; k < N / 2; k++) {
            odd[k] = x[2 * k + 1];
        }
        Complex[] r = fft(odd);

        // combine
        Complex[] y = new Complex[N];
        for (int k = 0; k < N / 2; k++) {
            double kth = -2 * k * Math.PI / N;
            Complex wk = new Complex(Math.cos(kth), Math.sin(kth));
            y[k] = q[k].add(wk.multiply(r[k]));
            y[k + N / 2] = q[k].subtract(wk.multiply(r[k]));
        }
        return y;
    }


    /**
     * 计算频谱
     *
     * @param data         采样数据
     * @param Fs           采样率
     * @param needFillZero 是否需要补零
     * @param flyLen       参与傅里叶变换的长度
     */
    public static SpectrumDo caculateSpectrum(double[] data, double Fs, boolean needFillZero, int flyLen) {
        if (data == null || data.length == 0) {
            return null;
        }
        int dataLen = data.length;
        //参与傅里叶变换的数组
        double[] partakeFLYArr = null;
        //需要补零
        if (needFillZero) {
            if (flyLen < dataLen) {
                return null;
            }
            partakeFLYArr = new double[flyLen];
            System.arraycopy(data, 0, partakeFLYArr, 0, dataLen);
            for (int i = dataLen; i < flyLen; i++) {
                partakeFLYArr[i] = 0;
            }
        } else {
            partakeFLYArr = data;
            flyLen = dataLen;
        }
        //进行傅里叶变换,使用的是commons-math包里面的快速傅里叶算法
        FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
        Complex[] complexArr = fft.transform(partakeFLYArr, TransformType.FORWARD);

        SpectrumDo spectrumBean = new SpectrumDo();
        double[] amplitude = new double[flyLen / 2];

        //计算振幅,除以的都是采样点的长度,而不是傅里叶变化的长度
        Complex complex = complexArr[0];
        amplitude[0] = Math.sqrt(Math.abs(complex.getReal()) * Math.abs(complex.getReal())
                + Math.abs(complex.getImaginary()) * Math.abs(complex.getImaginary())) * 1.0 / dataLen;
        for (int i = 1; i < flyLen / 2; i++) {
            complex = complexArr[i];
            amplitude[i] = Math.sqrt(Math.abs(complex.getReal()) * Math.abs(complex.getReal())
                    + Math.abs(complex.getImaginary()) * Math.abs(complex.getImaginary())) * 2.0 / dataLen;
        }
        //设置振幅
        spectrumBean.setAmplitude(amplitude);

        //设置频率分辨率
        spectrumBean.setFrequencyResolution(Fs / flyLen);

        return spectrumBean;
    }


    /**
     * Double转complex
     *
     * @param data
     * @return
     */
    public static Complex[] double2complex(Double[] data) {
        Complex[] complex = new Complex[data.length];
        //类数组需要二次初始化，切记，不然会造成空指针情况
        for (int index = 0; index < complex.length; index++) {
            complex[index] = new Complex(0, 0);
        }
        for (int index = 0; index < complex.length; index++) {
            //将时域的 y 传入实部即可，虚部默认为 0，注意这一步需要结合你的需要来做。
            complex[index] = new Complex(data[index], 0);
        }
        return complex;
    }


}
